package nsfs

import (
	"GLibHac/tools"
	"encoding/binary"
	"fmt"
	"github.com/pkg/errors"
)

type NcaCompressionInfo struct {
	MetaOffset int64
	MetaSize   int64
	MetaHeader *BucketTreeHeader
}

func NewNcaCompressionInfo(data []byte) *NcaCompressionInfo {
	return &NcaCompressionInfo{
		MetaOffset: int64(binary.LittleEndian.Uint64(data)),
		MetaSize:   int64(binary.LittleEndian.Uint64(data[8:])),
		MetaHeader: NewBucketTreeHeader(data[16:32]),
	}
}

type NcaFsPatchInfo struct {
	RelocationTreeOffset int64
	RelocationTreeSize   int64
	EncryptionTreeOffset int64
	EncryptionTreeSize   int64
}

func NewNcaFsPatchInfo(data []byte) *NcaFsPatchInfo {
	return &NcaFsPatchInfo{
		RelocationTreeOffset: int64(binary.LittleEndian.Uint64(data)),
		RelocationTreeSize:   int64(binary.LittleEndian.Uint64(data[0x08:])),
		EncryptionTreeOffset: int64(binary.LittleEndian.Uint64(data[0x20:])),
		EncryptionTreeSize:   int64(binary.LittleEndian.Uint64(data[0x28:])),
	}
}

type NcaSparseInfo struct {
	MetaOffset     int64
	MetaSize       int64
	MetaHeader     []byte
	PhysicalOffset int64
	Generation     uint64
}

func NewNcaSparseInfo(data []byte) (r *NcaSparseInfo) {
	r = &NcaSparseInfo{
		MetaOffset:     int64(binary.LittleEndian.Uint64(data)),
		MetaSize:       int64(binary.LittleEndian.Uint64(data[8:])),
		PhysicalOffset: int64(binary.LittleEndian.Uint64(data[32:])),
		Generation:     binary.LittleEndian.Uint64(data[40:]),
	}
	r.MetaHeader = data[16:32]
	return
}

const (
	Sha256LevelOffset = 0x28
	Sha256LevelSize   = 0x10
	MasterHashOffset  = 0
	MasterHashSize    = 0x20
)

type Sha256Level struct {
	Offset int64
	Size   int64
}

func NewSha256Level(data []byte) *Sha256Level {
	return &Sha256Level{
		Offset: int64(binary.LittleEndian.Uint64(data)),
		Size:   int64(binary.LittleEndian.Uint64(data[8:])),
	}
}

type NcaFsIntegrityInfoSha256 struct {
	data       []byte
	BlockSize  int32
	LevelCount int32
}

func (s *NcaFsIntegrityInfoSha256) ValidateLevelIndex(index int32) error {
	if index < 0 || index > 5 {
		return errors.Errorf("IVFC level index must be between 0 and 5. Actual: %d", index)
	}
	return nil
}

func (s *NcaFsIntegrityInfoSha256) GetLevelInfo(index int32) (*Sha256Level, error) {
	if err := s.ValidateLevelIndex(index); err != nil {
		return nil, err
	}
	offset := Sha256LevelOffset + Sha256LevelSize*index
	return NewSha256Level(s.data[offset:]), nil
}

func (s *NcaFsIntegrityInfoSha256) GetLevelOffset(index int32) (int64, error) {
	info, err := s.GetLevelInfo(index)
	if err != nil {
		return 0, err
	}
	return info.Offset, nil
}

func (s *NcaFsIntegrityInfoSha256) GetLevelSize(index int32) (int64, error) {
	info, err := s.GetLevelInfo(index)
	if err != nil {
		return 0, err
	}
	return info.Size, nil
}

func (s *NcaFsIntegrityInfoSha256) GetMasterHash() []byte {
	return s.data[MasterHashOffset : MasterHashOffset+MasterHashSize]
}

func NewNcaFsIntegrityInfoSha256(data []byte) (r *NcaFsIntegrityInfoSha256) {
	r = &NcaFsIntegrityInfoSha256{
		data:       data,
		BlockSize:  tools.BytesToIntLittleEndian(data[0x20:]),
		LevelCount: tools.BytesToIntLittleEndian(data[0x24:]),
	}
	return
}

const (
	IvfcLevelsOffset = 0x10
	IvfcLevelSize    = 0x18
)

type IvfcLevel struct {
	Offset    int64
	Size      int64
	BlockSize int32
}

func NewIvfcLevel(data []byte) *IvfcLevel {
	return &IvfcLevel{
		Offset:    int64(binary.LittleEndian.Uint64(data)),
		Size:      int64(binary.LittleEndian.Uint64(data[8:])),
		BlockSize: tools.BytesToIntLittleEndian(data[0x10:]),
	}
}

type NcaFsIntegrityInfoIvfc struct {
	data           []byte
	Magic          uint32
	Version        int32
	MasterHashSize int32
	LevelCount     int32
}

func (s *NcaFsIntegrityInfoIvfc) ValidateLevelIndex(index int32) error {
	if index < 0 || index > 6 {
		return errors.Errorf("IVFC level index must be between 0 and 6. Actual: %d", index)
	}
	return nil
}

func (s *NcaFsIntegrityInfoIvfc) GetMasterHash() []byte {
	return s.data[MasterHashOffset : MasterHashOffset+s.MasterHashSize]
}

func (s *NcaFsIntegrityInfoIvfc) GetLevelInfo(index int32) (*IvfcLevel, error) {
	if err := s.ValidateLevelIndex(index); err != nil {
		return nil, err
	}
	offset := IvfcLevelsOffset + IvfcLevelSize*index
	return NewIvfcLevel(s.data[offset:]), nil
}

func (s *NcaFsIntegrityInfoIvfc) GetLevelOffset(index int32) (int64, error) {
	info, err := s.GetLevelInfo(index)
	if err != nil {
		return 0, err
	}
	return info.Offset, nil
}

func (s *NcaFsIntegrityInfoIvfc) GetLevelSize(index int32) (int64, error) {
	info, err := s.GetLevelInfo(index)
	if err != nil {
		return 0, err
	}
	return info.Size, nil
}

func NewNcaFsIntegrityInfoIvfc(data []byte) *NcaFsIntegrityInfoIvfc {
	return &NcaFsIntegrityInfoIvfc{
		data:           data,
		Magic:          binary.LittleEndian.Uint32(data),
		Version:        tools.BytesToIntLittleEndian(data[4:]),
		MasterHashSize: tools.BytesToIntLittleEndian(data[8:]),
		LevelCount:     tools.BytesToIntLittleEndian(data[12:]),
	}
}

type NcaHashType byte

func (h NcaHashType) String() string {
	switch h {
	case NcaHashTypeAuto:
		return "Auto"
	case NcaHashTypeNone:
		return "None"
	case NcaHashTypeSha256:
		return "Sha256"
	case NcaHashTypeIvfc:
		return "Ivfc"
	default:
		return fmt.Sprintf("NcaHashType(%d)", h)
	}
}

const (
	NcaHashTypeAuto = NcaHashType(iota)
	NcaHashTypeNone
	NcaHashTypeSha256
	NcaHashTypeIvfc
)

type NcaFormatType byte

func (f NcaFormatType) String() string {
	switch f {
	case NcaFormatTypeRomfs:
		return "Romfs"
	case NcaFormatTypePfs0:
		return "Pfs0"
	default:
		return fmt.Sprintf("NcaFormatType(%d)", f)
	}
}

const (
	NcaFormatTypeRomfs = NcaFormatType(iota)
	NcaFormatTypePfs0
)

type NcaEncryptionType byte

func (f NcaEncryptionType) String() string {
	switch f {
	case NcaEncryptionTypeAuto:
		return "Auto"
	case NcaEncryptionTypeNone:
		return "None"
	case NcaEncryptionTypeXTS:
		return "XTS"
	case NcaEncryptionTypeAesCtr:
		return "AesCtr"
	case NcaEncryptionTypeAesCtrEx:
		return "AesCtrEx"
	default:
		return fmt.Sprintf("NcaEncryptionType(%d)", f)
	}
}

const (
	NcaEncryptionTypeAuto = NcaEncryptionType(iota)
	NcaEncryptionTypeNone
	NcaEncryptionTypeXTS
	NcaEncryptionTypeAesCtr
	NcaEncryptionTypeAesCtrEx
)

type NcaFsHeader struct {
	HeaderData      []byte
	CompressionInfo *NcaCompressionInfo
	Version         int16
	FormatType      NcaFormatType
	HashType        NcaHashType
	EncryptionType  NcaEncryptionType
	UpperCounter    uint64
	CounterType     uint32
	CounterVersion  uint32
	PatchInfo       *NcaFsPatchInfo
	*NcaSparseInfo
	*NcaFsIntegrityInfoSha256
	*NcaFsIntegrityInfoIvfc
}

func (h *NcaFsHeader) IsPatchSection() bool {
	return h.PatchInfo.RelocationTreeSize != 0
}

func (h *NcaFsHeader) ExistsSparseLayer() bool {
	return h.NcaSparseInfo.Generation != 0
}

func (h *NcaFsHeader) ExistsCompressionLayer() bool {
	return h.CompressionInfo.MetaOffset != 0 && h.CompressionInfo.MetaSize != 0
}

func NewNcaFsHeader(headerData []byte) *NcaFsHeader {
	return &NcaFsHeader{
		CompressionInfo:          NewNcaCompressionInfo(headerData[0x178 : 0x178+0x20]),
		Version:                  int16(binary.LittleEndian.Uint16(headerData)),
		FormatType:               NcaFormatType(headerData[2]),
		HashType:                 NcaHashType(headerData[3]),
		EncryptionType:           NcaEncryptionType(headerData[4]),
		UpperCounter:             binary.LittleEndian.Uint64(headerData[0x140:]),
		CounterType:              binary.LittleEndian.Uint32(headerData[0x140:]),
		CounterVersion:           binary.LittleEndian.Uint32(headerData[0x144:]),
		PatchInfo:                NewNcaFsPatchInfo(headerData[0x100 : 0x100+0x40]),
		NcaSparseInfo:            NewNcaSparseInfo(headerData[0x148 : 0x148+0x30]),
		NcaFsIntegrityInfoSha256: NewNcaFsIntegrityInfoSha256(headerData[8 : 8+0xF8]),
		NcaFsIntegrityInfoIvfc:   NewNcaFsIntegrityInfoIvfc(headerData[8 : 8+0xF8]),
	}
}
